package com.nebula.core.client;

import com.nebula.core.jdbc.NebulaResultSet;
import com.zaxxer.hikari.pool.HikariProxyResultSet;
import lombok.Data;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.DataSourceUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.Statement;

/**
 * @author zhangtao
 */
@Data
public class NebulaSession {

    private DataSource dataSource;

    private static final Logger log = LoggerFactory.getLogger(NebulaSession.class);

    public NebulaSession(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    public com.vesoft.nebula.client.graph.data.ResultSet query(String ngql) {
        log.info("执行ngql：{}", ngql);
        NebulaResultSet nebulaResultSet = null;
        try (Connection connection = dataSource.getConnection();
             Statement statement = connection.createStatement();
             HikariProxyResultSet hikariProxyResultSet = (HikariProxyResultSet) statement.executeQuery(ngql)) {
            nebulaResultSet = hikariProxyResultSet.unwrap(NebulaResultSet.class);
            DataSourceUtils.releaseConnection(connection, dataSource);
        } catch (Exception e) {
            log.error(String.format("Execute: `%s', failed: %s",
                    ngql, e.getMessage()), e);
            throw new RuntimeException(e.getMessage());
        }
        return nebulaResultSet.getResultSet();
    }

    public boolean execute(String ngql) {
        log.info("执行ngql：{}", ngql);
        try (Connection connection = dataSource.getConnection();
             Statement statement = connection.createStatement();) {
            boolean result = statement.execute(ngql);
            DataSourceUtils.releaseConnection(connection, dataSource);
            return result;
        } catch (Exception e) {
            log.error("nebula语句执行出错 {}", e.getMessage());
            log.error("ngql : {}", ngql);
            throw new RuntimeException(e.getMessage());
        }
    }

}
